import torch.nn as nn

class FasterRCNNTrainer(nn.Module):
    def __init__(self, args):
        super(FasterRCNNTrainer, self).__init__()
    def forward(self, model_train, loss_fn, imgs, bboxes, labels):
        bs = imgs.shape[0]
        img_size = imgs.shape[2:]

        base_feature = model_train(imgs, mode = 'backbone')

        rpn_locs, rpn_scores, rois, roi_indices, anchor = model_train([base_feature, img_size],  mode = 'rpn_head')

        rpn_loc_loss_all, rpn_cls_loss_all, sample_rois, sample_indexes, gt_roi_locs, gt_roi_labels = loss_fn([rpn_locs, rpn_scores, rois, roi_indices, anchor], imgs, bboxes, labels, mode = 'rpn_head' )

        roi_cls_locs, roi_scores = model_train([base_feature, sample_rois, sample_indexes, img_size], mode = 'roi_head')

        roi_loc_loss_all, roi_cls_loss_all = loss_fn([roi_cls_locs, roi_scores, gt_roi_locs, gt_roi_labels], imgs, bboxes, labels, mode = 'roi_head')

        losses = [rpn_loc_loss_all/bs, rpn_cls_loss_all/bs, roi_loc_loss_all/bs, roi_cls_loss_all/bs]

        # losses的最后一项一定是'total_loss'!!!
        losses = losses + [sum(losses)]
        # 以字典形式返回
        losses_name = ['rpn_loc_loss','rpn_cls_loss','roi_loc_loss','roi_cls_loss','train_loss']
        losses = dict(zip(losses_name, losses))

        return losses

class YOLOv3Trainer(nn.Module):
    def __init__(self,args):
        super(YOLOv3Trainer, self).__init__()
    def forward(self, model_train, loss_fn, imgs, targets):
        outputs         = model_train(imgs)
        loss_value_all  = 0
        for l in range(len(outputs)):
            loss_item = loss_fn(l, outputs[l], targets)
            loss_value_all  += loss_item
        loss_value = loss_value_all  

        losses = loss_value

        losses_name = ['train_loss']
        losses = dict(zip(losses_name, losses))
        
        return losses
